from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, \
    VectorStoreIndex, ListIndex, ComposableGraph
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm
from llama_index.core.extractors.metadata_extractors import (
    KeywordExtractor,
    PydanticProgramExtractor,
    QuestionsAnsweredExtractor,
    SummaryExtractor,
    TitleExtractor,
)
from llama_index.core.extractors.document_context import DocumentContextExtractor



documents_part1=[Document(text="江苏的省会是南京")]

documents_part2=[Document(text='''南京，简称"宁"，别称"金陵""建康"，是江苏省省会、长三角核心城市，拥有3100年建城史和450年建都史，与北京、西安、洛阳并称中国四大古都‌
1。作为特大城市，其面积6587平方公里（相当于8个纽约曼哈顿），人口近千万，综合实力稳居全国前十‌
''')]


# 初始化各子索引
vector_index = VectorStoreIndex.from_documents(documents_part1)
keyword_index = ListIndex.from_documents(documents_part2)

# 构建组合图
graph = ComposableGraph.from_indices(
    root_index_cls=SummaryIndex,
    children_indices=[ vector_index,keyword_index],
    index_summaries=["江苏的省会", "南京的介绍"]
)

# 配置自定义检索策略
query_engine = graph.as_query_engine(
    vector_top_k=3,
    keyword_top_k=2
)

rs=query_engine.query("介绍一下江苏的省会")
print(rs)



